package com.hyw.trace.configuration.sample.aspect;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;

import com.hyw.trace.configuration.sample.annotation.JointPointType;
import com.hyw.trace.configuration.sample.annotation.MultiTraceSources;
import com.hyw.trace.configuration.sample.annotation.TraceSource;
import com.hyw.trace.configuration.sampler.AbstractSamplerTemplate;

@Aspect
public class MultiTraceSourcesAspect {

  private AbstractSamplerTemplate traceTemplate;

  public MultiTraceSourcesAspect(AbstractSamplerTemplate traceSamplerTemplate) {
    this.traceTemplate = traceSamplerTemplate;
  }

  @Around(value = "@annotation(multiTraceSources)")
  public Object around(ProceedingJoinPoint joinPoint, MultiTraceSources multiTraceSources) throws Throwable {
    Map<JointPointType, List<TraceSource>> mp = Arrays.asList(multiTraceSources.value()).stream()
        .collect(Collectors.groupingBy(TraceSource::jointPoint));

    Object ret = null;

    if (mp.containsKey(JointPointType.Before)) {
      mp.get(JointPointType.Before).forEach(s -> traceTemplate.run(joinPoint.getArgs(), s));
    }

    try {
      ret = joinPoint.proceed();
      return ret;
    } finally {
      if (mp.containsKey(JointPointType.After)) {
        mp.get(JointPointType.After).forEach(s -> traceTemplate.run(joinPoint.getArgs(), s));
      } else if (mp.containsKey(JointPointType.Result)) {
        for (TraceSource s : mp.get(JointPointType.Result)) {
          traceTemplate.run(new Object[] { ret }, s);
        }
      }
    }
  }

}
